import os
import sys

sys.path.append(os.path.dirname(os.path.dirname(__file__)))
from util_my import add_arg
from project.yolov5.data import refactor

def get_weight(git_dir='', cfg_id='', weights='', out_dir=''):

    weight_path = weights or os.path.join(out_dir, 'exp', 'weights', 'best.pt')

    if not os.path.exists(weight_path):
        try:
            from downloads.yolov5.utils import downloads
        except:
            sys.path.append(os.path.dirname(git_dir))
            from yolov5.utils import downloads

        weight_path = os.path.join(git_dir, cfg_id + '.pt')
        downloads.attempt_download(weight_path)

    return weight_path


def run(task='', data_dir='', cfg_id='', github_dir='', device='', weights='', out_dir='', 
    main_queue=None, sub_queue=None, exist=True, augment=False, im_sz='', logger=None,
    infer_dir='', conf_thres=0):

    git_dir, code_dir = refactor(task=task, data_dir=data_dir, github_dir=github_dir)

    is_cls = task == 'classification' or 'cls' in data_dir
    is_seg = task == 'segmentation' or 'seg' in data_dir
    is_det = task == 'detection' or 'det' in data_dir

    weights = get_weight(git_dir=git_dir, cfg_id=cfg_id, weights=weights, out_dir=out_dir)

    add_arg('--device', device)
    add_arg('--weights', weights)
    add_arg('--source', infer_dir)
    add_arg('--imgsz', im_sz or None)
    add_arg('--augment', augment or None)
    add_arg('--project', out_dir or None)
    add_arg('--exist-ok' if exist else None)
    add_arg('--conf-thres', conf_thres if not is_cls else None or None)
    add_arg(None if is_cls else '--nosave --save-txt --save-conf')

    root = os.path.dirname(__file__)
    while root in sys.path:
        sys.path.remove(root)
    while os.path.dirname(root) in sys.path:
        sys.path.remove(os.path.dirname(root))

    try:
        if is_cls:
            import downloads.yolov5.classify.predict as predict
        elif is_seg:
            import downloads.yolov5.segment.predict as predict
        elif is_det:
            import downloads.yolov5.detect as predict
    except:
        sys.path.append(os.path.dirname(git_dir))
        if is_cls:
            import yolov5.classify.predict as predict
        elif is_seg:
            import yolov5.segment.predict as predict
        elif is_det:
            import yolov5.detect as predict

    opt = predict.parse_opt()
    if 'development' not in code_dir:
        opt.main_ = main_queue
        opt.sub = sub_queue

    try:
        predict.main(opt)
    except Exception as e:
        if logger:
            logger.log(logger.ERROR, e, exception=True)
        if sub_queue:
            sub_queue.put({'err': str(e)})
